package cn.doitedu.sparkml.bayes

import java.io.File
import java.util

import cn.doitedu.commons.utils.{FileUtils, SparkUtil}
import com.hankcs.hanlp.HanLP
import com.hankcs.hanlp.seg.common.Term
import org.apache.log4j.{Level, Logger}
import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.feature.{HashingTF, IDF, Tokenizer}
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}

import scala.collection.mutable

/**
  * @author: 余辉
  * @blog: https://blog.csdn.net/silentwolfyh
  * @create: 2019/10/21
  * @description:
  * 需求：评论分类贝叶斯模型训练器
  *
  * 步鄹：
  * 1、spark读取测试数据，差评权重为1.0,一般权重为2.0,好评权重为3.0
  * 2、将差评，好评，中评通过差评进行合并，通过HanLP进行分词。（权重设置label，词变成特征）
  * 2、通过HashingTF类和IDF类预处理
  * 3、最后通过NaiveBayes类将label和doc的训练列传输
  * 4、保存训练结果
  **/
object BayesCommentModelTrainer {

  def main(args: Array[String]): Unit = {

    // 1、建立session连接
    Logger.getLogger("org").setLevel(Level.WARN)
    val spark: SparkSession = SparkUtil.getSparkSession(this.getClass.getSimpleName)
    import spark.implicits._

    val poor: Dataset[String] = spark.read.textFile("C:\\Work\\mldata\\comment\\poor")
    val eneral: Dataset[String] = spark.read.textFile("C:\\Work\\mldata\\comment\\general")
    val good: Dataset[String] = spark.read.textFile("C:\\Work\\mldata\\comment\\good")

    val poors: Dataset[(Double, String)] = poor.map(line => (1.0, line))
    val enerals: Dataset[(Double, String)] = eneral.map(line => (2.0, line))
    val goods: Dataset[(Double, String)] = good.map(line => (3.0, line))

    val resultData: Dataset[(Double, String)] = poors.union(enerals).union(goods)

    val labelAndFeatures: DataFrame = resultData.rdd.map(tp => {
      val label = tp._1
      import scala.collection.JavaConversions._
      val terms: util.List[Term] = HanLP.segment(tp._2)
      val features: mutable.Buffer[String] = terms.map(word => word.word).filter(_.length > 1)
      (label, features)
    }).toDF("label", "features")

    val tf: HashingTF = new HashingTF().setInputCol("features").setOutputCol("tf").setNumFeatures(1000000)
    val tfDF: DataFrame = tf.transform(labelAndFeatures)

    val idf: IDF = new IDF().setInputCol("tf").setOutputCol("idf")
    val idfDF: DataFrame = idf.fit(tfDF).transform(tfDF)

    val bayes: NaiveBayes = new NaiveBayes().setLabelCol("label").setFeaturesCol("idf").setSmoothing(1)
    val model: NaiveBayesModel = bayes.fit(idfDF)

    FileUtils.deleteDir(new File("rec_system/outputdata/comment_bayes_model"))
    model.write.save("rec_system/outputdata/comment_bayes_model")
  }
}
